from PIL import Image
import os
import numpy as np


PATH = r'testDigits'
NUM = (20, 10)
BORDER = 1  # pixel, width of border


def load_data(digit):
    border = BORDER
    lightness = 0.8
    f_list = os.listdir(PATH)

    img = np.zeros([NUM[1]*(32+border), NUM[0]*(32+border)])
    '''
    img[0] = lightness
    img[-1] = lightness
    img[:, 0] = lightness
    img[:, -1] = lightness
    for x in range(NUM[0]):
        img[:, x*(32+1)] = lightness

    for y in range(NUM[1]):
        img[y*(32+1)] = lightness
    '''
    for i in range(NUM[0]):
        for j in range(NUM[1]):
            for f_name in f_list:
                f_list.remove(f_name)
                if f_name.split('_')[0] == str(digit):
                    print(f_name, i, j)
                    img[j*32:(j+1)*32, i*32:(i+1)*32] = load_file(os.path.join(PATH, f_name))
                    break

    return img


def load_file(f_name):
    arr = np.zeros([32, 32], dtype=int)

    with open(f_name, 'r') as f:
        lines = f.readlines()
        for i in range(32):
            for j in range(32):
                arr[i][j] = int(lines[i][j])

    return arr


def visualize(digit):
    img = load_data(digit) * 255
    img = Image.fromarray(img)
    # img = img.convert(mode='L')

    return img


if __name__ == '__main__':
    for i in range(10):
        im = visualize(i)
        im = im.convert('L')
        im.save('test-%s.png' % str(i))

